"""
SG (Savitzky-Golay平滑) 处理器

SG滤波器用于平滑光谱数据，去除噪声的同时保持光谱特征。
使用多项式拟合进行局部平滑。
"""

import pandas as pd
import numpy as np
from scipy.signal import savgol_filter
from ..base import PreprocessorInterface
from ..validation import validate_input_data, safe_copy_dataframe


class SGSmooth(PreprocessorInterface):
    """
    Savitzky-Golay平滑滤波器
    
    使用多项式拟合对光谱数据进行平滑处理，
    在去除噪声的同时保持光谱的形状特征。
    """
    
    def __init__(self, window_length=11, polyorder=2):
        """
        初始化SG滤波器参数
        
        Args:
            window_length (int): 滤波窗口长度，必须为奇数，默认11
            polyorder (int): 多项式阶数，默认2
        """
        # 确保窗口长度为奇数
        if window_length % 2 == 0:
            window_length += 1
        
        self.window_length = window_length
        self.polyorder = polyorder
    
    def process(self, data: pd.DataFrame) -> pd.DataFrame:
        """
        对光谱数据进行SG平滑
        
        Args:
            data (pd.DataFrame): 输入光谱数据
            
        Returns:
            pd.DataFrame: SG平滑后的数据
        """
        # 验证输入数据
        validate_input_data(data)
        
        # 创建数据副本
        result = safe_copy_dataframe(data)
        
        # 获取光谱数据列（跳过第一列ID列）
        spectral_columns = result.columns[1:]
        spectral_data = result[spectral_columns].values
        
        # 检查窗口长度是否合适
        n_points = spectral_data.shape[1]
        window_length = min(self.window_length, n_points)
        
        # 确保窗口长度为奇数且大于多项式阶数
        if window_length % 2 == 0:
            window_length -= 1
        window_length = max(window_length, self.polyorder + 1)
        
        # 对每个样本进行SG平滑
        smoothed_spectra = np.zeros_like(spectral_data)
        
        for i in range(spectral_data.shape[0]):
            spectrum = spectral_data[i, :]
            
            try:
                # 应用Savitzky-Golay滤波器
                smoothed_spectrum = savgol_filter(
                    spectrum, 
                    window_length=window_length, 
                    polyorder=self.polyorder,
                    mode='nearest'  # 边界处理方式
                )
                smoothed_spectra[i, :] = smoothed_spectrum
            except Exception as e:
                # 如果平滑失败，保持原光谱
                print(f"SG平滑失败，样本 {i}: {e}")
                smoothed_spectra[i, :] = spectrum
        
        # 将平滑后的数据放回DataFrame
        result[spectral_columns] = smoothed_spectra
        
        return result
    
    def get_name(self) -> str:
        """返回处理器名称"""
        return f"SG平滑(窗口{self.window_length},阶数{self.polyorder})"
    
    def get_id(self) -> int:
        """返回唯一ID"""
        return 6